import torch

class QwenAspectRatioSelectorLatent:
    """
    Un node tout-en-un qui sélectionne une résolution optimisée pour Qwen Image
    et génère directement le tenseur latent vide correspondant.
    Il remplace le node "Empty Latent Image".
    """
    # Dictionnaire des ratios d'aspect avec des noms descriptifs
    ASPECT_RATIOS = {
        "1:1 (Perfect Square)": (1, 1),
        "2:3 (Classic Portrait)": (2, 3),
        "3:4 (Golden Ratio)": (3, 4),
        "3:5 (Elegant Vertical)": (3, 5),
        "4:5 (Artistic Frame)": (4, 5),
        "5:7 (Balanced Portrait)": (5, 7),
        "5:8 (Tall Portrait)": (5, 8),
        "7:9 (Modern Portrait)": (7, 9),
        "9:16 (Slim Vertical)": (9, 16),
        "9:19 (Tall Slim)": (9, 19),
        "9:21 (Ultra Tall)": (9, 21),
        "9:32 (Skyline)": (9, 32),
        "3:2 (Golden Landscape)": (3, 2),
        "4:3 (Classic Landscape)": (4, 3),
        "5:3 (Wide Horizon)": (5, 3),
        "5:4 (Balanced Frame)": (5, 4),
        "7:5 (Elegant Landscape)": (7, 5),
        "8:5 (Cinematic View)": (8, 5),
        "9:7 (Artful Horizon)": (9, 7),
        "16:9 (Panorama)": (16, 9),
        "19:9 (Cinematic Ultrawide)": (19, 9),
        "21:9 (Epic Ultrawide)": (21, 9),
        "32:9 (Extreme Ultrawide)": (32, 9),
    }

    # Liste des résolutions recommandées pour Qwen Image
    RECOMMENDED_RESOLUTIONS = [
        "1328x1328", "1664x928", "928x1664",
        "1472x1140", "1140x1472", "1584x1056",
        "1056x1584", "1024x1024", "768x768", "512x512",
    ]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mode": (["Recommandé", "Ratio d'Aspect", "Manuel"],),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
                "recommended": (s.RECOMMENDED_RESOLUTIONS,),
                "aspect_ratio": (list(s.ASPECT_RATIOS.keys()),),
                "target_size": ("INT", {
                    "default": 1328, "min": 256, "max": 4096, "step": 8, "display": "number"
                }),
                "manual_width": ("INT", {
                    "default": 1024, "min": 64, "max": 8192, "step": 8, "display": "number"
                }),
                "manual_height": ("INT", {
                    "default": 1024, "min": 64, "max": 8192, "step": 8, "display": "number"
                }),
            }
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    FUNCTION = "generate"
    CATEGORY = "Latent" # Catégorie changée pour être avec les autres nodes de Latent

    def generate(self, mode, batch_size, recommended, aspect_ratio, target_size, manual_width, manual_height):
        final_width, final_height = 1024, 1024

        if mode == "Manuel":
            final_width = manual_width
            final_height = manual_height

        elif mode == "Recommandé":
            try:
                parts = recommended.split('x')
                final_width = int(parts[0])
                final_height = int(parts[1])
            except:
                print("Erreur : La résolution recommandée est mal formatée. Utilisation des valeurs par défaut.")
                final_width, final_height = 1024, 1024
        
        elif mode == "Ratio d'Aspect":
            w_ratio, h_ratio = self.ASPECT_RATIOS[aspect_ratio]
            
            if w_ratio > h_ratio:  # Paysage
                final_width = target_size
                final_height = (target_size * h_ratio) / w_ratio
            elif h_ratio > w_ratio:  # Portrait
                final_height = target_size
                final_width = (target_size * w_ratio) / h_ratio
            else:  # Carré
                final_width = target_size
                final_height = target_size

            # Arrondir à un multiple de 8 pour la compatibilité VAE
            final_width = int(round(final_width / 8.0)) * 8
            final_height = int(round(final_height / 8.0)) * 8

        # Création du tenseur latent (divisé par 8 car c'est l'espace latent)
        latent_width = final_width // 8
        latent_height = final_height // 8
        
        # 4 canaux pour les modèles de diffusion standards
        latent = torch.zeros([batch_size, 4, latent_height, latent_width])

        return ({"samples": latent}, )

# Dictionnaire pour enregistrer le node dans ComfyUI
NODE_CLASS_MAPPINGS = {
    "QwenAspectRatioSelectorLatent": QwenAspectRatioSelectorLatent
}

# Nom d'affichage du node dans le menu de ComfyUI
NODE_DISPLAY_NAME_MAPPINGS = {
    "QwenAspectRatioSelectorLatent": "Qwen Aspect Ratio Selector (Latent)"
}